import pysoem
import struct
from ctypes import *
import time
import numpy as np
import mujoco
import threading

# 导入自定义函数和CANFD SDK动态链接库>>>
from os.path import dirname, join
import sys
import platform
arch = platform.machine()
sys.path.append(join(dirname(__file__), 'src/lysdemo_'+arch))
sys.path.append(join(dirname(__file__), "src"))
sys.path.append(dirname(__file__))
from src import *

site_force = np.zeros(6)

def xjc():
    global site_force
    # 创建 EtherCAT 主站，打开指定网口（请确保 enp4s0 已正确配置）
    master = pysoem.Master()
    master.open('enp4s0')

    # 扫描从站设备
    if master.config_init() == 0:
        print("未发现任何从站设备")
        master.close()
        return

    sensor = master.slaves[0]
    print("发现传感器:", sensor.name)

    # 进入 PRE-OP 状态，此状态允许通过 SDO 访问设备对象字典
    master.state = pysoem.PREOP_STATE
    master.write_state()
    # 稍作延时确保状态切换成功
    time.sleep(0.1)

    print("开始持续读取 SDO 数据，按 Ctrl+C 停止...")

    # 循环读取每个通道数据
    indices = [0x6000, 0x6001, 0x6002, 0x6003, 0x6004, 0x6005]
    channel_names = ['Fx', 'Fy', 'Fz', 'Mx', 'My', 'Mz']
    force_sensor = np.zeros(6)
    try:
        while True:
            # readings = {}
            for idx, name in zip(indices, channel_names):
                try:
                    # 读取对象 0x6000.1 ~ 0x6005.1
                    data = sensor.sdo_read(idx, 1)
                    # 根据文档，该数据类型为 REAL (32位 float)，使用小端格式解码
                    force_sensor[idx - 24576] = struct.unpack("<f", data)[0]
                    # value = struct.unpack("<f", data)[0]
                    # readings[name] = value
                    # print(f"idx {idx}, name {name}, value {value}")
                except Exception as e:
                    force_sensor[idx - 24576] = 0.0
                    print(f"Error: {e}")
                    # readings[name] = f"Error: {e}"
            # print("SDO 读取结果:", readings)
            # 对于7自由度，目前不需要变换
            # site_force = sensor_to_site(force_sensor)
            site_force = force_sensor
            time.sleep(0.001)  # 每 1ms 读取一次，可根据需要调整
    except KeyboardInterrupt:
        print("读取中断，程序退出")
    finally:
        master.close()


def mujoco_show():
    import mujoco.viewer as viewer
    pos_sensor_names = [
        "Motor_Shoulder_pos", 
        "Motor_BigArm1_pos",
        "Motor_BigArm2_pos",
        "Motor_SmallArm1_pos",
        "Motor_SmallArm2_pos",
        "Motor_Hand_pos",
        "Motor_Finger_pos"
        ]

    vel_sensor_names = [
        "Motor_Shoulder_vel", 
        "Motor_BigArm1_vel",
        "Motor_BigArm2_vel",
        "Motor_SmallArm1_vel",
        "Motor_SmallArm2_vel",
        "Motor_Hand_vel",
        "Motor_Finger_vel"
        ]


    # mujoco加载模型
    model_7DOF = mujoco.MjModel.from_xml_path(join(dirname(__file__),'scene_7DOF.xml')) 
    data_7DOF = mujoco.MjData(model_7DOF)
    mocap_id_7DOF = model_7DOF.body("target").mocapid[0]

    # TODO 导纳控制系数
    # TODO site force 转换到 base frame

    with viewer.launch_passive(model_7DOF,data_7DOF) as viewer:
        # Initialize the camera view to that of the free camera.
        mujoco.mjv_defaultFreeCamera(model_7DOF, viewer.cam)
        viewer.opt.frame = mujoco.mjtFrame.mjFRAME_SITE

        while viewer.is_running():

            print(f"site force is {site_force}")
            q_real_float = np.array([data_7DOF.sensor(pos_sensor_name).data[0] for pos_sensor_name in pos_sensor_names])
            qv_real_float = np.array([data_7DOF.sensor(vel_sensor_name).data[0] for vel_sensor_name in vel_sensor_names])

            mocap_pos = data_7DOF.mocap_pos[mocap_id_7DOF]

            mocap_quat = data_7DOF.mocap_quat[mocap_id_7DOF]  

            data_7DOF.ctrl = ik7(q_real_float, mocap_pos, mocap_quat, qv_real_float, site_force)# 控制
            # data.qpos = ik7(q_real_float, mocap_pos, mocap_quat, qv_real_float, site_force)# 展示ik位置
            mujoco.mj_step(model_7DOF, data_7DOF)
            mujoco.mj_forward(model_7DOF, data_7DOF)

            viewer.sync()
            time.sleep(model_7DOF.opt.timestep)


if __name__ == "__main__":
    sensor_thread = threading.Thread(target=xjc)
    sensor_thread.start()
    
    # MuJoCo在主线程运行（图形渲染推荐）
    mujoco_show()
    
    sensor_thread.join()